library(tidyverse)
library(ggplot2)
library(ggrepel)
library(readr)
library(lme4)
library(interactions)
library(patchwork)
library(ggridges)
library(jtools)
library(modelsummary)
library(colorspace)
library(interplot)
library(ggpubr)
library(margins)
library(estimatr)
library(sandwich)
library(patchwork)
library(stargazer)

# Load the data
data <- read_dta("Study_1.dta")

# Create treatment variable
data <- data %>%
  mutate(treatment = case_when(
    trustExperimentSplit == 1 ~ 0,
    trustExperimentSplit == 2 ~ 1
  )) %>%
  mutate(treatment = factor(treatment, labels = c("Control", "Treatment")))

# Create trust variables
data <- data %>%
  mutate(
    trustMP = coalesce(trustExperimentGroup1MPRe, trustExperimentGroup2MPRe),
    trustMSP = coalesce(trustExperimentGroup1MSPRe, trustExperimentGroup2MSPRe),
    trustScotMin = coalesce(trustExperimentGroup1ScotMinRe, trustExperimentGroup2ScotMinRe),
    trustUKMin = coalesce(trustExperimentGroup1UKMinRe, trustExperimentGroup2UKMinRe),
    trustCivWest = coalesce(trustExperimentGroup1CivWestRe, trustExperimentGroup2CivWestRe),
    trustCivScot = coalesce(trustExperimentGroup1CivScotRe, trustExperimentGroup2CivScotRe)
  )

# Create binary incumbent_con variable
data <- data %>% 
  mutate(incumbent_con = case_when(
    voteIntWestminster == 1 ~ 1,
    TRUE ~ 0
  ))

# Create subsets for Conservative voters and non-Conservative voters
con_voters <- subset(data, incumbent_con == 1)
non_con_voters <- subset(data, incumbent_con == 0)

# Function to run models and create plots for each outcome
create_outcome_plot <- function(outcome_var, data_con, data_non_con) {
  # Determine overall range for y-axis
  y_min <- min(c(data_con[[outcome_var]], data_non_con[[outcome_var]]), na.rm = TRUE)
  y_max <- max(c(data_con[[outcome_var]], data_non_con[[outcome_var]]), na.rm = TRUE)
  
  # Calculate the range and add a small buffer (e.g., 5% of the range)
  y_range <- y_max - y_min
  buffer <- 0.10 * y_range
  
  y_limits <- c(max(0, y_min - buffer), min(5, y_max + buffer))
  
  # Models
  model_con <- lm(as.formula(paste(outcome_var, "~ treatment")), data = data_con)
  model_non_con <- lm(as.formula(paste(outcome_var, "~ treatment")), data = data_non_con)
  
  # Predictions
  data_con$predicted <- predict(model_con, data_con)
  data_non_con$predicted <- predict(model_non_con, data_non_con)
  
  # Subsets for treatment and control
  treat_con <- subset(data_con, treatment == "Treatment")
  control_con <- subset(data_con, treatment == "Control")
  treat_non_con <- subset(data_non_con, treatment == "Treatment")
  control_non_con <- subset(data_non_con, treatment == "Control")
  
  # Calculate ATE and p-values
  ate_con <- coef(model_con)["treatmentTreatment"]
  ate_non_con <- coef(model_non_con)["treatmentTreatment"]
  p_value_con <- summary(model_con)$coefficients["treatmentTreatment", "Pr(>|t|)"]
  p_value_non_con <- summary(model_non_con)$coefficients["treatmentTreatment", "Pr(>|t|)"]
  
  # Function to add asterisks based on p-value
  add_asterisks <- function(p_value) {
    if (p_value < 0.01) return("***")
    else if (p_value < 0.05) return("**")
    else if (p_value < 0.1) return("*")
    else return("")
  }
  
  # Create plots
  plot_con <- effect_plot(model = model_con, pred = treatment, robust = TRUE,
                          cat.geom = "point", cat.interval.geom = "linerange",
                          colors = "black", cat.pred.point.size = 3, int.width = .90) +
    labs(title = paste("Conservative Voters (N =", nrow(data_con), ")")) +
    ylab(paste("Trust in", gsub("trust", "", outcome_var))) +
    xlab("") +
    scale_x_discrete(labels = c("Control" = "Control", "Treatment" = "Treatment")) +
    scale_y_continuous(limits = y_limits) +
    geom_jitter(data = treat_con, aes(x = treatment, y = .data[[outcome_var]]),
                height = .25, size = 3, width = .35, alpha = .15, shape = 20,
                pch = 21, color = "#11CE9E") +
    geom_jitter(data = control_con, aes(x = treatment, y = .data[[outcome_var]]),
                height = .25, width = .35, alpha = .15, shape = 17,
                pch = 21, size = 3, color = "#CE1141FF") +  
    geom_bracket(xmin = "Control", xmax = "Treatment",
                 y.position = y_limits[2] * 0.95,
                 label = sprintf("ATE=%.2f%s", ate_con, add_asterisks(p_value_con)),
                 tip.length = 0.02,
                 color = "black") +
    theme_minimal(base_size = 14) +
    theme(axis.title.y = element_text(face = "bold"),
          axis.text.x = element_text(face = "bold"))
  
  plot_non_con <- effect_plot(model = model_non_con, pred = treatment, robust = TRUE,
                              cat.geom = "point", cat.interval.geom = "linerange",
                              colors = "black", cat.pred.point.size = 3, int.width = .90) +
    labs(title = paste("Non-Conservative Voters (N =", nrow(data_non_con), ")")) +
    ylab("") +
    xlab("") +
    scale_x_discrete(labels = c("Control" = "Control", "Treatment" = "Treatment")) +
    scale_y_continuous(limits = y_limits) +
    geom_jitter(data = treat_non_con, aes(x = treatment, y = .data[[outcome_var]]),
                height = .25, width = .35, alpha = .15, shape = 20,
                pch = 21, size = 3, color = "#11CE9E") +
    geom_jitter(data = control_non_con, aes(x = treatment, y = .data[[outcome_var]]),
                height = .25, width = .35, alpha = .15, shape = 17,
                pch = 21, size = 3, color = "#CE1141FF") +
    geom_bracket(xmin = "Control", xmax = "Treatment",
                 y.position = y_limits[2] * 0.95,
                 label = sprintf("ATE=%.2f%s", ate_non_con, add_asterisks(p_value_non_con)),
                 tip.length = 0.02,
                 color = "black") +
    theme_minimal(base_size = 14) +
    theme(axis.title.y = element_blank(),
          axis.text.y = element_blank(),
          axis.text.x = element_text(face = "bold"))
  
  # Combine plots and return
  plot_con + plot_non_con +
    plot_annotation(title = paste('Effect of treatment on trust in', gsub("trust", "", outcome_var)),
                    theme = theme(plot.title = element_text(size = 14, face = "bold")))
}

# Create plots for each outcome
outcomes <- c("trustMP", "trustMSP", "trustUKMin", "trustScotMin", "trustCivWest", "trustCivScot")
plots <- map(outcomes, ~create_outcome_plot(.x, con_voters, non_con_voters))

# Combine all plots into one
combined_plot <- wrap_plots(plots, ncol = 2) +
  plot_annotation(
    title = 'Effects of treatment on trust in various entities',
    caption = "Treatment group outcome statistically distinct at p<.1(*), p<0.05(**), & p<0.01(***)",
    theme = theme(plot.title = element_text(size = 20, face = "bold"))
  )

# Save the combined plot
ggsave("Figure_8.png", combined_plot, width = 20, height = 30, dpi = 300)

# Create models
models <- list()
for (outcome in outcomes) {
  models[[outcome]] <- list(
    all = lm(as.formula(paste(outcome, "~ treatment")), data = data),
    interaction = lm(as.formula(paste(outcome, "~ treatment * incumbent_con")), data = data),
    con = lm(as.formula(paste(outcome, "~ treatment")), data = con_voters),
    non_con = lm(as.formula(paste(outcome, "~ treatment")), data = non_con_voters)
  )
}

# Create and save tables
for (outcome in outcomes) {
  stargazer(models[[outcome]], 
            type = "latex",
            title = paste("Models for", outcome),
            column.labels = c("All", "Interaction", "Conservative", "Non-Conservative"),
            out = paste0("Table_", outcome, "_incumbent_con.tex"))
}